import math
from typing import Any


def serialize_complex(c: complex) -> dict[str, float]:
    """将复数序列化为 JSON 可序列化的字典"""
    return {"real": c.real, "imag": c.imag}


def serialize_su2(su2: list[list[complex]] | None) -> list[list[dict[str, float]]] | None:
    """将 SU(2) 矩阵（包含复数）序列化为 JSON 可序列化的格式"""
    if su2 is None:
        return None
    return [[serialize_complex(c) for c in row] for row in su2]


def matrix_trace(m: list[list[float]]) -> float:
    return float(m[0][0] + m[1][1] + m[2][2])


def matrix_det(m: list[list[float]]) -> float:
    # 3x3 determinant
    return (
        m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
        - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
        + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0])
    )


def safe_arccos(x: float) -> float:
    return math.acos(max(-1.0, min(1.0, x)))


def axis_angle_from_o3(m: list[list[float]]) -> tuple[list[float] | None, float | None]:
    """
    Compute axis (unit vector) and angle (radians) from a proper rotation matrix.
    Returns (axis, angle). For identity, returns ([0,0,1], 0.0).
    For improper operations (det < 0), returns (None, None).
    """
    det = matrix_det(m)
    if det < 0.0:
        return None, None

    tr = matrix_trace(m)
    angle = safe_arccos((tr - 1.0) / 2.0)
    # Handle near-identity and pi cases carefully
    if abs(angle) < 1e-10:
        return [0.0, 0.0, 1.0], 0.0

    sin_a = math.sin(angle)
    if abs(sin_a) < 1e-10:
        # angle ~ pi; axis from diagonal elements
        # Find the axis using eigenvector approach approximation
        x = math.sqrt(max(0.0, (m[0][0] + 1.0) / 2.0))
        y = math.sqrt(max(0.0, (m[1][1] + 1.0) / 2.0))
        z = math.sqrt(max(0.0, (m[2][2] + 1.0) / 2.0))
        # choose signs using off-diagonals
        if m[0][1] < 0.0:
            y = -y
        if m[0][2] < 0.0:
            z = -z
        norm = math.sqrt(x * x + y * y + z * z) or 1.0
        return [x / norm, y / norm, z / norm], angle

    x = (m[2][1] - m[1][2]) / (2.0 * sin_a)
    y = (m[0][2] - m[2][0]) / (2.0 * sin_a)
    z = (m[1][0] - m[0][1]) / (2.0 * sin_a)
    norm = math.sqrt(x * x + y * y + z * z) or 1.0
    return [x / norm, y / norm, z / norm], angle


def su2_from_o3(m: list[list[float]]) -> list[list[complex]] | None:
    """
    Map a proper rotation in O(3) (det=+1) to SU(2) via the standard homomorphism.
    For det=-1 (improper operation), returns None.
    """
    axis, angle = axis_angle_from_o3(m)
    if axis is None or angle is None:
        return None

    nx, ny, nz = axis
    c = math.cos(angle / 2.0)
    s = math.sin(angle / 2.0)
    i = 1j
    return [
        [c - i * nz * s, (-i * nx - ny) * s],
        [(-i * nx + ny) * s, c + i * nz * s],
    ]


def parse_character_value(char_str: str | int | float) -> complex:
    """
    解析特征标表中的字符串值为复数

    支持的格式：
    - 整数/浮点数：1, -1, 0.5
    - 纯虚数：i, -i, 2i
    - 指数形式：e^(2πi/3), e^{2\\pi i/3}
    """
    # 如果已经是数字，直接返回
    if isinstance(char_str, (int, float)):
        return complex(char_str, 0)

    s = str(char_str).strip()

    # 处理 ω (omega, 三次单位根 e^(2πi/3))
    # ω = e^(2πi/3) = -1/2 + sqrt(3)/2 * i
    # ω^2 = e^(4πi/3) = -1/2 - sqrt(3)/2 * i
    omega = complex(-0.5, math.sqrt(3) / 2)  # e^(2πi/3)
    omega2 = complex(-0.5, -math.sqrt(3) / 2)  # e^(4πi/3)

    if s == "ω" or s == "omega":
        return omega
    if s == "-ω" or s == "-omega":
        return -omega
    if s == "ω^2" or s == "omega^2":
        return omega2
    if s == "-ω^2" or s == "-omega^2":
        return -omega2

    # 处理纯虚数 i, -i, 2i 等
    if s == "i":
        return 1j
    if s == "-i":
        return -1j
    if "i" in s and not "e" in s.lower() and not "π" in s and not "pi" in s and not "ω" in s:
        # 形如 2i, -3i
        coef = s.replace("i", "").strip()
        if coef == "" or coef == "+":
            return 1j
        if coef == "-":
            return -1j
        return complex(0, float(coef))

    # 处理指数形式 e^(2πi/3) 或 e^{2\pi i/3}
    if "e^" in s or "e^{" in s:
        import re

        # 提取指数部分
        pattern = r"e\^\{?([^}]+)\}?"
        match = re.search(pattern, s)
        if match:
            exp_str = match.group(1)
            # 计算指数值（假设形如 2πi/3, 4πi/3, πi 等）
            exp_value = eval_exponential(exp_str)
            return complex(math.cos(exp_value), math.sin(exp_value))

    # 默认尝试转换为浮点数
    try:
        return complex(float(s), 0)
    except:
        # 无法解析，返回 0
        return complex(0, 0)


def eval_exponential(exp_str: str) -> float:
    """
    计算指数表达式的值（假设形如 2πi/3, 4πi/3 等）
    返回实际的角度值（弧度）
    """
    import re

    # 移除空格和 i
    s = exp_str.replace(" ", "").replace("i", "").replace("\\", "")

    # 替换 π 和 pi 为实际值
    s = s.replace("π", f"*{math.pi}")
    s = s.replace("pi", f"*{math.pi}")

    # 安全计算表达式
    try:
        return eval(s, {"__builtins__": {}}, {"pi": math.pi})
    except:
        return 0.0


def infer_irrep_dimension(irrep: dict[str, Any]) -> int:
    """
    从 irrep 数据推断维度

    Args:
        irrep: 不可约表示数据 {"gamma": "Γ1", "characters": [1, 1, ...], ...}

    Returns:
        维度（1, 2, 3, ...）
    """
    # 如果有显式 dimension 字段
    if "dimension" in irrep:
        return int(irrep["dimension"])

    # 否则从 E 类（第一个）的特征标推断
    characters = irrep.get("characters", [])
    if characters and isinstance(characters[0], (int, float)):
        return int(abs(characters[0]))

    # 默认为 1D
    return 1


def generate_1d_irrep_matrices(
    character_table: dict[str, Any], operations: list[dict[str, Any]], classes_map: dict[str, list[str]] | None
) -> dict[int, dict[str, list[list[dict[str, float]]]]]:
    """
    为所有 1D 不可约表示生成矩阵

    Args:
        character_table: 特征标表 {"classes": [...], "irreps": [...]}
        operations: 操作列表，每个操作包含 id 和 seitz
        classes_map: 类映射 {"C1": ["1"], "C2": ["2_001"], ...}

    Returns:
        操作ID -> irrep标签 -> 1×1矩阵（序列化为 [[{"real": ..., "imag": ...}]]）
    """
    if not character_table or not classes_map:
        return {}

    classes = character_table.get("classes", [])
    irreps = character_table.get("irreps", [])

    # 构建 seitz -> 类索引 映射
    # classes_map 的 keys 是有序的（C1, C2, C3, ...），与特征标表的列顺序对应
    seitz_to_class_idx: dict[str, int] = {}

    # classes_map 的 keys 应该按顺序对应特征标表的列
    # 例如：classes_map = {"C1": ["1"], "C2": ["3+_001"], ...}
    #      对应 classes = ["E", "C3", ...] 的第 0, 1, ... 列
    for idx, (class_label, seitz_list) in enumerate(classes_map.items()):
        for seitz in seitz_list:
            # 清理 seitz 符号（移除 $ 和空格）
            clean_seitz = seitz.replace("$", "").replace(" ", "")
            seitz_to_class_idx[clean_seitz] = idx

    # 为每个操作生成 irrep 矩阵
    result: dict[int, dict[str, list[list[dict[str, float]]]]] = {}

    for op in operations:
        op_id = op.get("id")
        if op_id is None:
            continue

        seitz = op.get("seitz", "").replace(" ", "")

        # 查找该操作对应的类索引
        class_idx = seitz_to_class_idx.get(seitz)
        if class_idx is None:
            continue

        # 为每个 1D irrep 生成矩阵
        op_irreps: dict[str, list[list[dict[str, float]]]] = {}

        for irrep in irreps:
            # 只处理 1D irreps
            dim = infer_irrep_dimension(irrep)
            if dim != 1:
                continue

            irrep_label = irrep.get("gamma", irrep.get("mulliken", ""))
            characters = irrep.get("characters", [])

            if class_idx < len(characters):
                char_value = parse_character_value(characters[class_idx])
                # 1D irrep 矩阵是 1×1，即标量
                matrix_1x1 = [[serialize_complex(char_value)]]
                op_irreps[irrep_label] = matrix_1x1

        if op_irreps:
            result[op_id] = op_irreps

    return result


def generate_2d_irrep_from_o3(mat3x3: list[list[float]]) -> list[list[dict[str, float]]]:
    """
    从单个 O(3) 矩阵生成 2D irrep 矩阵（取 xy 平面投影）

    Args:
        mat3x3: 3×3 O(3) 矩阵

    Returns:
        2×2 矩阵（序列化为复数格式）
    """
    # 取 xy 平面的 2x2 子矩阵
    mat2x2 = [[mat3x3[0][0], mat3x3[0][1]], [mat3x3[1][0], mat3x3[1][1]]]

    # 序列化为复数格式
    return [[serialize_complex(complex(mat2x2[i][j], 0)) for j in range(2)] for i in range(2)]


def generate_all_irrep_matrices(
    character_table: dict[str, Any], operations: list[dict[str, Any]], classes_map: dict[str, list[str]] | None = None
) -> dict[int, dict[str, list[list[dict[str, float]]]]]:
    """
    为所有不可约表示生成矩阵（1D, 2D, 3D）

    Args:
        character_table: 特征标表 {"classes": [...], "irreps": [...], "class_operations": [[1], [2,3], ...]}
        operations: 操作列表，每个操作包含 id, seitz, matrix
        classes_map: (已废弃) 类映射 {"C1": ["1"], "C2": ["2_001"], ...}，仅用于向后兼容

    Returns:
        操作ID -> irrep标签 -> 矩阵
    """
    if not character_table:
        return {}

    irreps = character_table.get("irreps", [])

    # 优先使用新格式：class_operations (operation ids)
    class_operations = character_table.get("class_operations")

    if class_operations:
        # 新格式：直接使用 operation id 映射
        id_to_class_idx: dict[int, int] = {}
        for class_idx, op_ids in enumerate(class_operations):
            for op_id in op_ids:
                id_to_class_idx[op_id] = class_idx
    elif classes_map:
        # 旧格式：使用 seitz 符号映射（向后兼容）
        id_to_class_idx = {}
        seitz_to_class_idx: dict[str, int] = {}
        for idx, (class_label, seitz_list) in enumerate(classes_map.items()):
            for seitz in seitz_list:
                clean_seitz = seitz.replace("$", "").replace(" ", "")
                seitz_to_class_idx[clean_seitz] = idx

        # 构建 id -> class_idx 映射
        for op in operations:
            seitz = op.get("seitz", "").replace(" ", "")
            class_idx = seitz_to_class_idx.get(seitz)
            if class_idx is not None:
                id_to_class_idx[op["id"]] = class_idx
    else:
        return {}

    # 为每个操作生成 irrep 矩阵
    result: dict[int, dict[str, list[list[dict[str, float]]]]] = {}

    for op in operations:
        op_id = op.get("id")
        if op_id is None:
            continue

        matrix = op.get("matrix")

        # 查找该操作对应的类索引
        class_idx = id_to_class_idx.get(op_id)
        if class_idx is None or matrix is None:
            continue

        # 为每个 irrep 生成矩阵
        op_irreps: dict[str, list[list[dict[str, float]]]] = {}

        for irrep in irreps:
            irrep_label = irrep.get("gamma", irrep.get("mulliken", ""))
            characters = irrep.get("characters", [])
            dim = infer_irrep_dimension(irrep)

            if class_idx >= len(characters):
                continue

            if dim == 1:
                # 1D irrep: 特征标直接作为矩阵
                char_value = parse_character_value(characters[class_idx])
                matrix_1x1 = [[serialize_complex(char_value)]]
                op_irreps[irrep_label] = matrix_1x1

            elif dim == 2:
                # 2D irrep: 从 O(3) 矩阵投影到 xy 平面
                matrix_2x2 = generate_2d_irrep_from_o3(matrix)
                op_irreps[irrep_label] = matrix_2x2

            elif dim == 3:
                # 3D irrep: 直接使用 O(3) 矩阵
                matrix_3x3 = [[serialize_complex(complex(matrix[i][j], 0)) for j in range(3)] for i in range(3)]
                op_irreps[irrep_label] = matrix_3x3

        if op_irreps:
            result[op_id] = op_irreps

    return result


def summarize_operation(m: list[list[float]]) -> dict[str, Any]:
    """
    Summarize a 3x3 O(3) matrix operation.
    Returns dict with det, axis, angle, and su2 (if proper rotation).
    SU2 matrices are serialized to JSON-compatible format.
    """
    det = matrix_det(m)
    axis, angle = axis_angle_from_o3(m)
    su2_raw = su2_from_o3(m)

    # 序列化 SU(2) 矩阵中的复数
    su2_serialized = serialize_su2(su2_raw)

    return {
        "det": det,
        "axis": axis,
        "angle": angle,
        "su2": su2_serialized,
    }
